import sys
import numpy as np
from datetime import datetime

# tolerances for float16
relative_tol = 1e-3
absolute_tol = 1e-5
error_tol = 1e-3

def build_eye(rows, cols, k=0, dtype=np.float16):
    mat = np.zeros((rows, cols), dtype=dtype)
    for i in range(rows):
        j = i + k
        if 0 <= j < cols:
            mat[i, j] = 1.0
    return mat.reshape(-1)

def verify_eye(output_path, rows, cols, k=0, dtype_str="float16"):
    try:
        dtype = np.float16 if dtype_str == "float16" else np.float32
        output = np.fromfile(output_path, dtype=dtype)
        expected_size = rows * cols
        if output.size != expected_size:
            print(f"[ERROR] size mismatch: expected {expected_size}, got {output.size}")
            return False
        golden = build_eye(rows, cols, k, dtype=dtype)

        different = ~np.isclose(output, golden, rtol=relative_tol, atol=absolute_tol, equal_nan=True)
        idxs = np.where(different)[0]

        print("=== Eye Verification ===")
        if idxs.size == 0:
            print("All elements PASS")
        else:
            for i in idxs[:10]:
                g = golden[i]
                o = output[i]
                rdiff = abs(o - g) / g if g != 0 else (abs(o) if g == 0 else float('inf'))
                print(f"index: {i:06d}, expected: {g:.6f}, actual: {o:.6f}, rdiff: {rdiff:.6f}")
        error_ratio = idxs.size / golden.size
        print(f"\nerror ratio: {error_ratio:.6f}, tolerance: {error_tol:.6f}")
        return error_ratio <= error_tol
    except Exception as e:
        print(f"[ERROR] Verification failed: {e}")
        return False

if __name__ == "__main__":
    # usage:
    # python verify_result.py output.bin rows cols [k] [dtype]
    exit_code = 0
    try:
        if len(sys.argv) < 4:
            print("Usage: python verify_result.py output.bin rows cols [k=0] [dtype=float16]")
            sys.exit(1)
        output_path = sys.argv[1]
        rows = int(sys.argv[2])
        cols = int(sys.argv[3])
        k = int(sys.argv[4]) if len(sys.argv) > 4 else 0
        dtype_str = sys.argv[5] if len(sys.argv) > 5 else "float16"
        res = verify_eye(output_path, rows, cols, k, dtype_str)
        if not res:
            print("[ERROR] result error")
            exit_code = 1
        else:
            print("test pass")
    except Exception as e:
        print(e)
        exit_code = 1
    finally:
        print("timestamp:", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
        sys.exit(exit_code)
